import pandas as pd
from preprocess import bs_preprocess
from txt2json import RecordDataJson,Frame
import os
from bs37ToARKit import ARKitBlendShapeLocation
import numpy as np
from enum import Enum
import json
from dataclasses import dataclass,asdict, field

class ARKitBlendShapeLocationUpper(Enum):
    BROWDOWNLEFT = 1
    BROWDOWNRIGHT = 2
    BROWINNERUP = 3
    BROWOUTERUPLEFT = 4
    BROWOUTERUPRIGHT = 5
    CHEEKPUFF = 6
    CHEEKSQUINTLEFT = 7
    CHEEKSQUINTRIGHT = 8
    EYEBLINKLEFT = 9
    EYEBLINKRIGHT = 10
    EYELOOKDOWNLEFT = 11
    EYELOOKDOWNRIGHT = 12
    EYELOOKINLEFT = 13
    EYELOOKINRIGHT = 14
    EYELOOKOUTLEFT = 15
    EYELOOKOUTRIGHT = 16
    EYELOOKUPLEFT = 17
    EYELOOKUPRIGHT = 18
    EYESQUINTLEFT = 19
    EYESQUINTRIGHT = 20
    EYEWIDELEFT = 21
    EYEWIDERIGHT = 22
    JAWFORWARD = 23
    JAWLEFT = 24
    JAWOPEN = 25
    JAWRIGHT = 26
    MOUTHCLOSE = 27 
    MOUTHDIMPLELEFT = 28
    MOUTHDIMPLERIGHT = 29
    MOUTHFROWNLEFT = 30
    MOUTHFROWNRIGHT = 31
    MOUTHFUNNEL = 32
    MOUTHLEFT = 33
    MOUTHLOWERDOWNLEFT = 34
    MOUTHLOWERDOWNRIGHT = 35
    MOUTHPRESSLEFT = 36
    MOUTHPRESSRIGHT = 37
    MOUTHPUCKER = 38
    MOUTHRIGHT = 39
    MOUTHROLLLOWER = 40
    MOUTHROLLUPPER = 41
    MOUTHSHRUGLOWER = 42
    MOUTHSHRUGUPPER = 43
    MOUTHSMILELEFT = 44
    MOUTHSMILERIGHT = 45
    MOUTHSTRETCHLEFT = 46
    MOUTHSTRETCHRIGHT = 47
    MOUTHUPPERUPLEFT = 48
    MOUTHUPPERUPRIGHT = 49
    NOSESNEERLEFT = 50
    NOSESNEERRIGHT = 51
    TONGUEOUT = 52
    MAX = 53


def csvToUnityBSJson(srcCsvPath,bs_scale=1.0,bs_delay=0.0):
    
    rdj = RecordDataJson()
    time = 0
    step = 1.0 / 220

    csv_data = pd.read_csv(srcCsvPath)
    for index, row in csv_data.iterrows():
        frame = Frame()
        frame.time = time + bs_delay
        time = time + step
        frame.location = list(range(52))

        inputLen = len(csv_data.columns)
        weights = np.zeros(ARKitBlendShapeLocationUpper.MAX.value-2)
        
        
        for key, _ in csv_data.items():
            if key == 'time':
                frame.time = row[key]
                
        for key, _ in csv_data.items():
            ukey = key.upper()
            if any(x for x in ARKitBlendShapeLocationUpper if x.name == ukey):
                loc = ARKitBlendShapeLocationUpper[ukey].value-1
                v = row[key]
                weights[loc] = v * 100 * bs_scale
            
        weights = weights.tolist()

        weights.append(0.0)#last tounge bs
        frame.weight = weights
        rdj.frames.append(frame)
    
    
    jstr = json.dumps(asdict(rdj))
    print(jstr)
    
    jfilePath = srcCsvPath + '.bs_weight.json'
    f2 = open(jfilePath, 'w')
    f2.write(jstr)
    f2.close()
if __name__ == '__main__':

 
    dstPath = 'D:/DevelopProj/Yuji/Audio2Face/YAudio2BS/data/lenovo/csv/' 
    
    # srcCsvPath = os.path.join('data/lenovo_bs.csv')
    
    for root, _, files in os.walk(dstPath):
            for file in files:
                if file.endswith('.csv'):
                    srcCsvPath = os.path.join(root,file)    
    
                    csvToUnityBSJson(srcCsvPath)

